"""
STS加密模块
实现DES加密解密功能，用于保护数据传输安全
"""

import struct
from typing import Optional, Tuple
from Crypto.Cipher import DES
from Crypto.Util.Padding import pad, unpad


class STSCrypto:
    """STS加密解密类"""
    
    def __init__(self, key: Optional[bytes] = None):
        """
        初始化加密模块
        
        Args:
            key: 8字节DES密钥，如果为None则使用默认密钥
        """
        if key is None:
            # 默认密钥（实际应用中应该从安全存储中获取）
            self.key = b'\x01\x02\x03\x04\x05\x06\x07\x08'
        else:
            if len(key) != 8:
                raise ValueError("DES密钥必须是8字节")
            self.key = key
    
    def encrypt(self, data: bytes) -> bytes:
        """
        DES加密
        
        Args:
            data: 要加密的数据
            
        Returns:
            加密后的数据
        """
        cipher = DES.new(self.key, DES.MODE_ECB)
        padded_data = pad(data, DES.block_size)
        return cipher.encrypt(padded_data)
    
    def decrypt(self, encrypted_data: bytes) -> bytes:
        """
        DES解密
        
        Args:
            encrypted_data: 加密的数据
            
        Returns:
            解密后的数据
        """
        cipher = DES.new(self.key, DES.MODE_ECB)
        decrypted_data = cipher.decrypt(encrypted_data)
        return unpad(decrypted_data, DES.block_size)
    
    def encrypt_token_data(self, meter_id: str, amount: int, 
                          transfer_counter: int, crc: int) -> bytes:
        """
        加密令牌数据
        
        Args:
            meter_id: 表计ID
            amount: 充值金额（以分为单位）
            transfer_counter: 传输计数器
            crc: 校验码
            
        Returns:
            加密后的令牌数据
        """
        # 构建令牌数据结构
        # 格式: [meter_id(4字节)][amount(3字节)][transfer_counter(2字节)][crc(2字节)]
        meter_id_bytes = int(meter_id).to_bytes(4, 'big')
        amount_bytes = amount.to_bytes(3, 'big')
        transfer_counter_bytes = transfer_counter.to_bytes(2, 'big')
        crc_bytes = crc.to_bytes(2, 'big')
        
        token_data = meter_id_bytes + amount_bytes + transfer_counter_bytes + crc_bytes
        return self.encrypt(token_data)
    
    def decrypt_token_data(self, encrypted_data: bytes) -> Tuple[str, int, int, int]:
        """
        解密令牌数据
        
        Args:
            encrypted_data: 加密的令牌数据
            
        Returns:
            (meter_id, amount, transfer_counter, crc)
        """
        decrypted_data = self.decrypt(encrypted_data)
        
        # 解析数据
        meter_id = int.from_bytes(decrypted_data[0:4], 'big')
        amount = int.from_bytes(decrypted_data[4:7], 'big')
        transfer_counter = int.from_bytes(decrypted_data[7:9], 'big')
        crc = int.from_bytes(decrypted_data[9:11], 'big')
        
        return str(meter_id), amount, transfer_counter, crc
    
    def calculate_crc16(self, data: bytes) -> int:
        """
        计算CRC16校验码
        
        Args:
            data: 要计算校验码的数据
            
        Returns:
            CRC16校验码
        """
        crc = 0xFFFF
        for byte in data:
            crc ^= byte
            for _ in range(8):
                if crc & 0x0001:
                    crc = (crc >> 1) ^ 0xA001
                else:
                    crc >>= 1
        return crc
    
    def verify_crc(self, data: bytes, expected_crc: int) -> bool:
        """
        验证CRC校验码
        
        Args:
            data: 数据
            expected_crc: 期望的CRC校验码
            
        Returns:
            校验是否通过
        """
        calculated_crc = self.calculate_crc16(data)
        return calculated_crc == expected_crc
